import cv2
import os
from PIL import Image
import argparse
from pathlib import Path
import torch
from config import get_config
from Learner import face_learner
from utils.utils import load_facebank, draw_box_name, prepare_facebank,draw_box
from retinaface import Retinaface
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='for face verification')
    parser.add_argument("-f", "--file_name", help="video file name",default='video.mp4', type=str)
    parser.add_argument("-s", "--save_name", help="output file name",default='recording', type=str)
    parser.add_argument('-th','--threshold',help='threshold to decide identical faces',default=1.54, type=float)
    parser.add_argument("-u", "--update", help="whether perform update the facebank",action="store_true")
    parser.add_argument("-tta", "--tta", help="whether test time augmentation",action="store_true")
    parser.add_argument("-c", "--score", help="whether show the confidence score",action="store_true")
    parser.add_argument("-b", "--begin", help="from when to start detection(in seconds)", default=0, type=int)
    parser.add_argument("-d", "--duration", help="perform detection for how long(in seconds)", default=0, type=int)
    
    args = parser.parse_args()
    
    conf = get_config(False)
    # 加载Retinaface
    retinaface = Retinaface()
    print('retinaface loaded')

    # 识别人脸的模型加载（insight_face）
    learner = face_learner(conf, True)
    learner.threshold = args.threshold
    if conf.device.type == 'cpu':
        learner.load_state(conf, 'cpu_final.pth', True, True)
    else:
        # 加载模型
        learner.load_state(conf, 'final.pth', True, True)
    learner.model.eval()
    print('learner loaded')
    
    if args.update:
        targets, names = prepare_facebank(conf, learner.model, retinaface, tta = args.tta)
        print('facebank updated')
    else:
        targets, names = load_facebank(conf)
        print('facebank loaded')
    

    path = os.path.join(os.path.dirname(__file__),'./data/facebank/ysx/ysx_1.png')
    # draw = cv2.imread(path)
    # draw = cv2.cvtColor(draw,cv2.COLOR_BGR2RGB)
    # draw = Image.fromarray(draw)
    
    while True:
        img = input('请输入图片名：')
        draw = cv2.imread(img)
        old_img = draw
        # frame = np.array(old_image)
        # # RGBtoBGR满足opencv显示格式
        # frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
        if draw is None:
            print('错误！请重新输入!')
            continue
        else:
            draw = cv2.cvtColor(draw,cv2.COLOR_BGR2RGB)
            draw = Image.fromarray(draw)
            
            bboxes, faces = retinaface.align_multi(draw, conf.face_limit, 16)

            # bboxes = bboxes[:,:-1] #shape:[10,4],only keep 10 highest possibiity faces
            bboxes = bboxes.astype(int)
            # bboxes = bboxes + [-1,-1,1,1] # personal choice
            results, score, dist = learner.infer(conf, faces, targets, True)
            print("结果为：", results)
            print(dist,score,bboxes)
            for idx,bbox in enumerate(bboxes):
                if args.score:
                    # cv2.rectangle(old_img, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
                    frame = draw_box_name(bbox, names[results[idx] + 1] + '_{:.2f}'.format(score[idx]), old_img)
                    print(bbox, names[results[idx] + 1] + '_{:.2f}'.format(score[idx]))
                else:
                    # frame = draw_box_name(bbox, names[results[idx] + 1], old_img)
                    print(bbox, names[results[idx] + 1])
                    cv2.rectangle(old_img,(bbox[0],bbox[1]),(bbox[2],bbox[3]),(0,0,255),6)
                    cv2.putText(old_img,names[results[idx] + 1],(bbox[0],bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 2,(0,255,0),3,cv2.LINE_AA)
            # cv2.cvResizeWindow("待拼接图像1", 800, 400);  
            cv2.namedWindow("Video",0);
            cv2.resizeWindow("Video", 640, 480);
            # cv2.namedWindow("Video",1);   
            cv2.imshow('Video', old_img)
            cv2.waitKey(0)